"""
Author: xubing
Date: 2024-04-28 02:17:53
LastEditors: xubing
LastEditTime: 2024-04-28 09:22:59
Description: file content
"""

import os

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import (
    TextLoader,
    UnstructuredPDFLoader,
    UnstructuredWordDocumentLoader,
)
from loguru import logger

from core.pdf_loader import UnstructuredPaddlePDFLoader


class DocPreprocessing:
    def __init__(self) -> None:
        """
        这里的处理逻辑是先使用loader
        再使用spliter
        """
        logger.info("|-> 开始处理原始数据...")

    def _load_document(self, file):
        file_type = file.lower().split(".")[-1]

        if file_type == "pdf":
            loader = UnstructuredPaddlePDFLoader(file)
            # loader = UnstructuredPDFLoader(file)
        elif file_type in ["docx", "doc"]:
            loader = UnstructuredWordDocumentLoader(file)
        elif file_type == "txt":
            loader = TextLoader(file)
        else:
            logger.error("不支持的格式")
            return None
        try:
            origin_docs = loader.load()
        except Exception as e:
            print(e)
        return origin_docs

    def process(self, file, split=False):
        origin_docs = self._load_document(file)
        if not split:
            # 这里返回原始文档,主要是为了高级rag, 构造父子文档块
            return origin_docs
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=100,
            chunk_overlap=10,
        )
        splited_docs = text_splitter.split_documents(origin_docs)
        logger.info("|-> 原始数据处理完毕!")
        return splited_docs

    def batch_process(self, file_list, split=False):
        if len(file_list) == 0:
            logger.error("File list can not be empty")
        docs = []
        for file in file_list:
            if not os.path.isfile(file):
                logger.warning(
                    f"The [{file}] is not a valid file, the process will continue..."
                )
            else:
                logger.info("Start process file [{}]".format(file))
                try:
                    # 处理文件
                    origin_docs = self.process(file, split=split)
                    docs.extend(origin_docs)
                except:
                    logger.error("文件处理失败,请检查!")
        return docs


if __name__ == "__main__":
    file_path = "data/new"
    file_list = [os.path.join(file_path, x) for x in os.listdir(file_path)]
    docs = DocPreprocessing().batch_process(file_list)
    print(docs)
